package cn.bincker.mybatis.encrypt.data.migration.impl;

import cn.bincker.mybatis.encrypt.data.migration.entity.MigrateTable;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.io.Closeable;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.BiFunction;
import java.util.stream.Collectors;

@Slf4j
public class SelectTask implements Runnable, Closeable {
    @Getter
    private boolean running = true;
    private boolean destroyed = false;
    private final DataSource dataSource;
    private MigrateTable table;
    private final BiFunction<MigrateTable, Long, String> querySqlBuilder;
    private final ExistQueryBuilder existQueryBuilder;
    private boolean checkExist;
    private final BlockingQueue<Long> pageQueue;
    private final Queue<Object[]> nextQueue;
    private final boolean failContinue;
    private final int pageSize;
    private final AtomicLong selectedRows;
    private final AtomicLong migratedSuccessRows;
    private final int nextQueueLimit;
    private final Object busyLock = new Object();
    private transient boolean busy = false;

    public SelectTask(DataSource dataSource, BiFunction<MigrateTable, Long, String> querySqlBuilder, ExistQueryBuilder existQueryBuilder, BlockingQueue<Long> pageQueue, Queue<Object[]> nextQueue, boolean failContinue, int pageSize, AtomicLong selectedRows, AtomicLong migratedSuccessRows, int nextQueueLimit) {
        this.dataSource = dataSource;
        this.querySqlBuilder = querySqlBuilder;
        this.existQueryBuilder = existQueryBuilder;
        this.pageQueue = pageQueue;
        this.nextQueue = nextQueue;
        this.failContinue = failContinue;
        this.pageSize = pageSize;
        this.selectedRows = selectedRows;
        this.migratedSuccessRows = migratedSuccessRows;
        this.nextQueueLimit = nextQueueLimit;
    }

    public void setMigrateTable(MigrateTable table, boolean checkExist){
        this.table = table;
        this.checkExist = checkExist;
    }

    @Override
    public void close() {
        running = false;
    }

    public void waitDestroyed(){
        if (destroyed) return;
        synchronized (this){
            try {
                this.wait();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public void run() {
        try(var conn = dataSource.getConnection(); var statement = conn.createStatement()) {
            while (running){
                Long page;
                page = pageQueue.poll(1, TimeUnit.SECONDS);
                if (page == null) {
                    synchronized (busyLock){
                        busy = false;
                        busyLock.notifyAll();
                    }
                    continue;
                }
                busy = true;
                String sql = querySqlBuilder.apply(table, page);
                try {
                    var result = statement.executeQuery(sql);
                    var rows = new ArrayList<Object[]>(pageSize);
                    while (result.next()) {
                        var row = new Object[table.getColumns().size()];
                        for (int i = 0; i < row.length; i++) {
                            row[i] = result.getObject(i + 1);
                        }
                        selectedRows.incrementAndGet();
                        rows.add(row);
                    }
                    if(checkExist) {
                        try (var existQuery = existQueryBuilder.invoke(conn, table, rows)) {
                            var existsRows = existQuery.executeQuery();
                            var rowMapByKey = rows.stream().collect(Collectors.toMap(
                                    row->table.getKeys().stream().map(k->row[k.getIndex()]).collect(Collectors.toList()),
                                    row->row
                            ));
                            while (existsRows.next()){
                                var key = new ArrayList<>(table.getKeys().size());
                                for (int j = 0; j < table.getKeys().size(); j++) {
                                    key.add(existsRows.getObject(j + 1));
                                }
                                rowMapByKey.remove(key);
                            }
                            int existRows = rows.size() - rowMapByKey.size();
                            table.getMigratedRows().getAndAdd(existRows);
                            migratedSuccessRows.getAndAdd(existRows);
                            if (rowMapByKey.isEmpty()) continue;
                            waitingForSpend();
                            nextQueue.addAll(rowMapByKey.values());
                        }
                    }else{
                        waitingForSpend();
                        nextQueue.addAll(rows);
                    }
                }catch (Throwable e){
                    if (failContinue){
                        log.error("query fail: sql={}", sql, e);
                    }else{
                        throw e;
                    }
                }
            }
        } catch (Throwable e) {
            log.error("select task error", e);
            running = false;
        }
        destroyed = true;
        synchronized (this) {
            this.notifyAll();
        }
    }

    private void waitingForSpend() throws InterruptedException {
        if (nextQueue.size() < nextQueueLimit) return;
        synchronized (nextQueue){
            while (nextQueue.size() > nextQueueLimit){
                nextQueue.wait(1000);
            }
        }
    }

    @FunctionalInterface
    public interface ExistQueryBuilder {
        PreparedStatement invoke(Connection conn, MigrateTable table, List<Object[]> rows) throws SQLException;
    }

    public void waitBusy(){
        while (busy) {
            synchronized (busyLock) {
                try {
                    busyLock.wait(1000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }
}
